package com.ajocer.springbootinit.manager;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.TimeUnit;

@Service
public class SparkAIManager extends WebSocketListener {
    private static final Logger log = LoggerFactory.getLogger(SparkAIManager.class);
    
    @Value("${xingfeng.api-host}")
    public String hostUrl;
    
    @Value("${xingfeng.app-id}")
    public String appid;
    
    @Value("${xingfeng.api-secret}")
    public String apiSecret;
    
    @Value("${xingfeng.api-key}")
    public String apiKey;

    private static final Gson gson = new Gson();
    public static String totalAnswer=""; // 大模型的答案汇总
    
    // 用户提问的问题
    public static String NewQuestion = "";

    private static String userId = UUID.randomUUID().toString().substring(0, 10);

    private static Boolean wsCloseFlag = false;  // 控制ws关闭

    private static Boolean totalFlag=true; // 控制提示用户是否输入
    // 构造函数
    public SparkAIManager(String userId, Boolean wsCloseFlag) {
        this.userId = userId;
        this.wsCloseFlag = wsCloseFlag;
    }

    // 主函数
    public String sendMessageAndGetResponse(String message,Integer sleepTime) throws Exception {
        NewQuestion = message;
        totalAnswer = "";
        wsCloseFlag = false;
        createConnection();
        while(!wsCloseFlag){
            Thread.sleep(sleepTime);
        }
        return totalAnswer;
    }

    public static boolean canAddHistory(){  // 由于历史记录最大上限2.2W左右，需要判断是能能加入历史
        int history_length=0;
        for(RoleContent temp:historyList){
            history_length+=temp.content.length();
        }
        if(history_length>22000){
            historyList.clear();
            return false;
        }else{
            return true;
        }
    }

    // 创建鉴权url
    public String getAuthUrl(String hostUrl, String apiKey, String apiSecret) throws Exception {
        URL url = new URL(hostUrl);
        // 时间
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());
        // 拼接
        String preStr = "host: " + url.getHost() + "\n" +
                "date: " + date + "\n" +
                "GET " + url.getPath() + " HTTP/1.1";
        // SHA256加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        // Base64加密
        String sha = Base64.getEncoder().encodeToString(hexDigits);
        // 拼接
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey, "hmac-sha256", "host date request-line", sha);
        // 拼接地址
        HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath())).newBuilder().//
                addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8))).//
                addQueryParameter("date", date).//
                addQueryParameter("host", url.getHost()).//
                build();

        return httpUrl.toString();
    }

    private static List<RoleContent> historyList=new ArrayList<>(); // 对话历史存储

    private void createConnection() throws Exception {
        // 获取鉴权url
        String authUrl = getAuthUrl(hostUrl, apiKey, apiSecret);
        OkHttpClient client = new OkHttpClient.Builder().build();
        String url = authUrl.replace("http://", "ws://").replace("https://", "wss://");
        Request request = new Request.Builder().url(url).build();
        // 最新问题
        RoleContent roleContent=new RoleContent();
        roleContent.role="user";
        roleContent.content=NewQuestion;
        historyList.add(roleContent);

        WebSocket webSocket = client.newWebSocket(request, this);
    }

    @Override
    public void onOpen(WebSocket webSocket, Response response) {
        super.onOpen(webSocket, response);
        // 发送消息
        JSONObject requestJson = new JSONObject();

        // 构建请求参数
        JSONObject header = new JSONObject();
        header.put("app_id", appid);
        header.put("uid", userId);

        JSONObject parameter = new JSONObject();
        JSONObject chat = new JSONObject();
        chat.put("domain", "x1");  // 更新为x1领域
        chat.put("temperature", 0.5);
        chat.put("max_tokens", 4096);
        parameter.put("chat", chat);

        JSONObject payload = new JSONObject();
        JSONObject message = new JSONObject();
        JSONArray text = new JSONArray();

        // 添加历史对话
        for (RoleContent tempContent:historyList){
            JSONObject temp = new JSONObject();
            temp.put("role",tempContent.role);
            temp.put("content",tempContent.content);
            text.add(temp);
        }

        message.put("text", text);
        payload.put("message", message);

        requestJson.put("header", header);
        requestJson.put("parameter", parameter);
        requestJson.put("payload", payload);
        webSocket.send(requestJson.toString());
    }

    @Override
    public void onMessage(WebSocket webSocket, String text) {
        // System.out.println(userId + "用来区分那个用户的结果" + text);
        JsonParse myJsonParse = gson.fromJson(text, JsonParse.class);

        if (myJsonParse.header.code != 0) {
            log.error("发生错误，错误码为：{}", myJsonParse.header.code);
            log.error("本次请求的sid为：{}", myJsonParse.header.sid);
            webSocket.close(1000, "");
        }

        List<Text> textList = myJsonParse.payload.choices.text;
        for (Text temp : textList) {
            log.debug("AI回答内容: {}", temp.content);
            totalAnswer += temp.content;
        }

        if (myJsonParse.header.status == 2) {
            // 可以关闭连接，释放资源
            log.info("*****************************");
            log.info("本次请求的sid为：{}", myJsonParse.header.sid);

            RoleContent roleContent=new RoleContent();
            roleContent.role="assistant";
            roleContent.content=totalAnswer;
            historyList.add(roleContent);
            canAddHistory();
            wsCloseFlag = true;
            webSocket.close(1000, "");
        }
    }

    @Override
    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
        super.onFailure(webSocket, t, response);
        try {
            if (null != response) {
                int code = response.code();
                log.error("onFailure code: {}", code);
                log.error("onFailure body: {}", response.body().string());
                if (101 != code) {
                    log.error("connection failed");
                    throw new RuntimeException("星火API连接失败");
                }
            }
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    // 返回的json结果拆解
    class JsonParse {
        Header header;
        Payload payload;
    }

    class Header {
        int code;
        int status;
        String sid;
    }

    class Payload {
        Choices choices;
    }

    class Choices {
        List<Text> text;
    }

    class Text {
        String role;
        String content;
    }
    class RoleContent{
        String role;
        String content;

        public String getRole() {
            return role;
        }

        public void setRole(String role) {
            this.role = role;
        }

        public String getContent() {
            return content;
        }

        public void setContent(String content) {
            this.content = content;
        }
    }
}
